from flask import Flask, render_template, jsonify, request
import os
from flask_sqlalchemy import SQLAlchemy
import appbuilder
from aip import AipImageClassify
from werkzeug.utils import secure_filename

app = Flask(__name__)

# 数据库配置（请根据你的 MySQL 配置修改）
app.config['SQLALCHEMY_DATABASE_URI'] = 'mysql+pymysql://root:root@localhost/b23015122db'
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
db = SQLAlchemy(app)

# 数据模型
class ImageQuery(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    image_path = db.Column(db.String(255), nullable=False)
    user_query = db.Column(db.String(255), nullable=False)
    answer = db.Column(db.String(255), nullable=False)

# 创建数据库表
with app.app_context():
    db.create_all()

class ImageAnalyzer:
    def __init__(self, app):
        self.app = app
        self.setup_config()
        self.setup_clients()
        self.register_routes()

    def setup_config(self):
        """设置应用配置，包括上传目录和文件大小限制"""
        self.app.config['UPLOAD_FOLDER'] = 'uploads'
        os.makedirs(self.app.config['UPLOAD_FOLDER'], exist_ok=True)
        self.app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 限制文件大小为16MB

    def setup_clients(self):
        """初始化API客户端"""
        os.environ["APPBUILDER_TOKEN"] = "bce-v3/ALTAK-X96iO5hyb2sN8KEhdiKoc/00e436f2edcb2733b26e19ef60f651c1aaa72fd3"
        app_id = "94a09af2-a754-4ace-aca0-ba899fc477cb"
        self.app_builder_client = appbuilder.AppBuilderClient(app_id)

        # 百度AI图像识别客户端
        APP_ID = '115878780'  # 替换为您的百度AI App ID
        API_KEY = '2ytPPKfWXLkk8xvRSK64mr0M'
        SECRET_KEY = 'CmdxptQFZDFTAgs539H3vOFy4sV61I0m'
        self.image_client = AipImageClassify(APP_ID, API_KEY, SECRET_KEY)

    def register_routes(self):
        """注册应用路由"""
        self.app.add_url_rule('/', 'index', self.index)
        self.app.add_url_rule('/analyze-image', 'analyze_image', self.analyze_image, methods=['POST'])
        self.app.add_url_rule('/ask-question', 'ask_question', self.ask_question, methods=['POST'])
        self.app.add_url_rule('/process-image-and-query', 'process_image_and_query', self.process_image_and_query, methods=['POST'])

    def index(self):
        """返回主页面的HTML"""
        return render_template('demo5.html')

    def analyze_image(self):
        """处理图像识别请求并返回结果"""
        return self.process_image(request.files['image'], self.image_client.advancedGeneral)

    def ask_question(self):
        """处理用户提问请求并返回AI的回答"""
        user_query = request.json.get('query')
        conversation_id = self.app_builder_client.create_conversation()
        response = self.app_builder_client.run(conversation_id, user_query)
        return jsonify({"answer": response.content.answer})

    def process_image_and_query(self):
        """结合图像和用户问题进行处理"""
        image_file = request.files['image']
        user_query = request.form.get('query')
        return self.process_image(image_file, self.image_client.advancedGeneral, user_query)

    def process_image(self, image_file, api_call, user_query=None):
        """处理图像文件并调用指定的API"""
        if image_file:
            filename = secure_filename(image_file.filename)
            file_path = os.path.join(self.app.config['UPLOAD_FOLDER'], filename)
            image_file.save(file_path)

            with open(file_path, 'rb') as f:
                image_data = f.read()

            options = {"baike_num": 10}
            image_result = api_call(image_data, options)

            image_description = (
                image_result['result'][0]['keyword']
                if "result" in image_result and image_result['result']
                else "无法识别图片内容"
            )

            if user_query:
                combined_query = f"图像描述: {image_description}. 用户的问题: {user_query}"
                conversation_id = self.app_builder_client.create_conversation()
                response = self.app_builder_client.run(conversation_id, combined_query)

                # 存储到数据库
                new_entry = ImageQuery(image_path=file_path, user_query=user_query, answer=response.content.answer)
                db.session.add(new_entry)
                db.session.commit()

                return jsonify({"answer": response.content.answer})

            return jsonify({"image_description": image_description})

        return jsonify({"error": "没有上传图片"}), 400

if __name__ == '__main__':
    analyzer = ImageAnalyzer(app)
    app.run(debug=True, port=5500)  # 启动Flask应用